use anyhow::Result;
use clap::Parser;
use rusqlite::{Connection, OpenFlags, };
use rustyline::error::ReadlineError;

#[derive(Parser)]
struct Args {
    #[arg(help = "Path to SQLite database file")]
    db_path: String,
}

fn main() -> Result<()> {
    let args = Args::parse();

    let flags: OpenFlags = OpenFlags::SQLITE_OPEN_READ_WRITE;
    let conn = match Connection::open_with_flags(&args.db_path, flags) {
        Ok(conn) => {
            println!("Connected to {}", args.db_path);
            conn
        },
        Err(e) => {
            anyhow::bail!("Failed to open database '{}': {}", args.db_path, e);
        }
    };

    println!("Type SQL commands or 'quit' to exit.");

    let mut rl = rustyline::DefaultEditor::new()?;
    loop {
        let readline = rl.readline("sqlite> ");
        match readline {
            Ok(line) => {
                let _ = rl.add_history_entry(line.as_str());
                if line.trim().eq_ignore_ascii_case("quit") {
                    break;
                }
                // 执行 SQL 并输出结果
                execute_sql(&conn, &line)?;
            }
            Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => break,
            Err(err) => {
                eprintln!("Readline error: {}", err);
                break;
            }
        }
    }

    conn.close().ok();
    Ok(())
}

fn execute_sql(conn: &Connection, sql: &str) -> Result<()> {
    // 尝试作为查询执行（SELECT等）
    match conn.prepare(sql) {
        Ok(mut stmt) => {
            // 获取列信息
            let column_count = stmt.column_count();
            if column_count == 0 {
                // 不返回结果的语句
                let rows_affected = stmt.execute([])?;
                println!("{} rows affected", rows_affected);
                return Ok(());
            }
            // 获取列名
            let column_names: Vec<String> = (0..column_count)
                .map(|i| stmt.column_name(i).unwrap_or("unknown").to_string())
                .collect();

            // 执行查询
            let mut rows = stmt.query([])?;

            // 打印表头
            println!("{}", column_names.join("\t"));
            println!("{}", "-".repeat(column_names.join("\t").len()));

            // 打印每一行数据
            while let Some(row) = rows.next()? {
                let values: Vec<String> = (0..column_count)
                    .map(|i| {
                        match row.get::<_, rusqlite::types::Value>(i) {
                            Ok(val) => format!("{:?}", val),
                            Err(_) => "NULL".to_string(),
                        }
                    })
                    .collect();
                println!("{}", values.join("\t"));
            }
        }
        Err(_) => {
            // 如果 prepare 失败，尝试作为非查询语句执行
            match conn.execute(sql, []) {
                Ok(rows_affected) => {
                    println!("{} rows affected", rows_affected);
                }
                Err(e) => {
                    eprintln!("Error: {}", e);
                }
            }
        }
    }
    Ok(())
}